#include "tmf_4d_inv.h"
#include <cuda.h>




/**
 * Perform the inversion of \f$ D_{dwf}\psi = \eta \f$ to find \f$ \psi \f$ where \f$ \psi \f$ and \f$ \eta \f$ are in Fourier space
 *
 * @param eta Source fermion
 * @param psi Solution Fermion
 * @param u Gauge field
 * @param M5 Mass M5
 * @param device Device number of GPU on which the inversion should be performed
 */
void tmf_4d_inv::twistedMassFermionInvert(LatticeFermion &eta, LatticeFermion &psi, multi1d<LatticeColorMatrix> &u,double M5, int device){
    ////////////////////////////////////////////////
    // QUDA PARAMETERS AND SETUP
    ////////////////////////////////////////////////
    // TODO: Better way of finding QDP precision
    QudaPrecision cpu_prec;
    switch ( sizeof(Real(0)) )
    {
        case 4:
            cpu_prec = QUDA_SINGLE_PRECISION;
            break;
        case 8:
            cpu_prec = QUDA_DOUBLE_PRECISION;
            break;
        default:
            std::cerr << "Unable to determine QDP++ precision.";
            exit(1);
    }
    // Using GTX 480, single precision
    QudaPrecision cuda_prec = QUDA_SINGLE_PRECISION;
    QudaPrecision cuda_prec_sloppy = QUDA_HALF_PRECISION;

    QudaGaugeParam gauge_param = newQudaGaugeParam();
    QudaInvertParam inv_param = newQudaInvertParam();

    multi1d<int> Lat = Layout::lattSize();
    gauge_param.X[0] = Lat[0];
    gauge_param.X[1] = Lat[1];
    gauge_param.X[2] = Lat[2];
    gauge_param.X[3] = Lat[3];

    gauge_param.anisotropy = 1.0;
    gauge_param.type = QUDA_WILSON_LINKS;
    gauge_param.gauge_order = QUDA_QDP_GAUGE_ORDER;
    gauge_param.t_boundary = QUDA_ANTI_PERIODIC_T;

    gauge_param.cpu_prec = cpu_prec;
    gauge_param.cuda_prec = cuda_prec;
    gauge_param.reconstruct = QUDA_RECONSTRUCT_12;
    gauge_param.cuda_prec_sloppy = cuda_prec_sloppy;
    gauge_param.reconstruct_sloppy = QUDA_RECONSTRUCT_12;
    gauge_param.gauge_fix = QUDA_GAUGE_FIXED_NO;

    inv_param.dslash_type = QUDA_TWISTED_MASS_DSLASH;
    inv_param.inv_type = QUDA_CG_INVERTER;
    inv_param.twist_flavor = QUDA_TWIST_PLUS;

    inv_param.tol = 5e-8;
    inv_param.maxiter = 10000;
    inv_param.reliable_delta = 1e-2;

    double mass = M5;
    inv_param.kappa = 1.0 / (2.0*(1 + 3/gauge_param.anisotropy + mass));
    inv_param.mu = 0.1;
    inv_param.twist_flavor = QUDA_TWIST_MINUS;

    inv_param.solution_type = QUDA_MAT_SOLUTION;
    inv_param.solve_type = QUDA_DIRECT_PC_SOLVE;
    inv_param.matpc_type = QUDA_MATPC_EVEN_EVEN;
    inv_param.dagger = QUDA_DAG_NO;
    inv_param.mass_normalization = QUDA_MASS_NORMALIZATION;

    inv_param.cpu_prec = cpu_prec;
    inv_param.cuda_prec = cuda_prec;
    inv_param.cuda_prec_sloppy = cuda_prec_sloppy;
    inv_param.preserve_source = QUDA_PRESERVE_SOURCE_YES;
    inv_param.dirac_order = QUDA_DIRAC_ORDER; // MAYBE CHANGE

    gauge_param.ga_pad = 0; // 24*24*24;
    inv_param.sp_pad = 0;   // 24*24*24;
    inv_param.cl_pad = 0;   // 24*24*24;

    inv_param.verbosity = QUDA_VERBOSE;

    ///////////////////////////////////////
    // End config and begin preparing data
    ///////////////////////////////////////
    
    void* gauge[4];
    for(int mu=0; mu < Nd; mu++)
    { 
        // Pointer to first element of Gauge field in each direction
        gauge[mu] = (void *)&(u[mu].elem(all.start()).elem().elem(0,0).real());
    }

    // initialize the QUDA library and import gauge field
    loadGaugeQuda(gauge, &gauge_param);

    //////////////////////////////////////////////////////////
    // Performing the inversion for each 'k'
    //////////////////////////////////////////////////////////

    QDPIO::cout << "#####################################\n";
    QDPIO::cout << "## Begin QUDA inversion";
    QDPIO::cout << "#####################################\n";
    QDPIO::cout << "~Source norm2 = " << norm2(eta) << " \n";
    void* spinorIn = (void *)&(eta.elem(all.start()).elem(0).elem(0).real());
    void* spinorOut =(void *)&((psi).elem(all.start()).elem(0).elem(0).real());
    // Perform the inversion
    invertQuda(spinorOut,spinorIn, &inv_param);
    // Print GPU usage if verbose
    QDPIO::cout << "~Solution norm2 = " << norm2(psi) << " \n";

    QDPIO::cout << "#####################################\n";
    QDPIO::cout << "## END QUDA INVERSIONS\n";
    QDPIO::cout << "#####################################\n";
    // finalize the QUDA library
}

